#! /usr/bin/env python
# -*- coding: utf-8 -*-

# Hipotetyczny rozszerzony algorytm verleta w 3D dla zadanego wektora siły F
# wg książki "Szczególna Teoria Względności" W.A. Ugarowa, str. 125, wzór 5.38

# Autorzy:
# Mateusz Madej
# Łukasz Rejmer
# Bartosz Rogalski
# Przemysław Piórkowski

try:
    import Gnuplot, Gnuplot.funcutils, numpy
except ImportError:
    pass

import math

def verlet(N, x0, xd, v0, delta, F):
    x1 = [x0, xd]; x2 = [x0, xd]; x3 = [x0, xd]
    v1 = [v0]; v2 = [v0]; v3 = [v0]
    X = [[0,0,0]]
    V = [[0,0,0]]
    
    c = 3 * math.pow(10,8)
   
    for i in range(1, N + 2):
        x1.append( 2 * x1[i] - x1[i - 1] + ( math.sqrt( 1 - math.pow( float( v1[i-1] / c ), 2 ))) * \
                 ( F[0] - ( v1[i-1] / math.pow( c,2 )) * ( F[0] * v1[i-1] )) * math.pow( delta,2 ))
        
        x2.append( 2 * x2[i] - x2[i - 1] + ( math.sqrt( 1 - math.pow( float( v2[i-1] / c ), 2 ))) * \
                 ( F[1] - ( v2[i-1] / math.pow( c,2 )) * ( F[1] * v2[i-1] )) * math.pow( delta,2 ))
        
        x3.append( 2 * x3[i] - x3[i - 1] + ( math.sqrt( 1 - math.pow( float( v3[i-1] / c ), 2 ))) * \
                 ( F[2] - ( v3[i-1] / math.pow( c,2 )) * ( F[2] * v3[i-1] )) * math.pow( delta,2 ))
        
        v1.append(( x1[i + 1] - x1[i - 1] ) / ( 2 * delta ))
        
        v2.append(( x2[i + 1] - x2[i - 1] ) / ( 2 * delta ))
        
        v3.append(( x3[i + 1] - x3[i - 1] ) / ( 2 * delta ))
        
        X.append([x1[i-1], x2[i-1], x3[i-1]])
        V.append([v1[i-1], v2[i-1], v3[i-1]])
    return X, V

def rysuj():
    N = 1000; x0 = 1; xd = 0; v0 = 1; delta = 0.01; F = [10, 330, 524]
    x, v = verlet(N, x0, xd, v0, delta, F)
    dane = [[],[],[]]    
    try:
        g = Gnuplot.Gnuplot(debug=1)
        for i in range(N):
            dane[0].append(x[i][0])
            dane[1].append(x[i][1])
            dane[2].append(x[i][2])
        dane = numpy.transpose((dane[0],dane[1],dane[2]))
        g.splot(Gnuplot.Data(dane, with_='linesp', inline=0))
        #g.hardcopy('wykres.ps', enhanced=1, color=1)
        raw_input(u'Naciśnij ENTER aby zakończyć...\n')
    except NameError:
        plik = open("dane.txt", "w")
        for i in range(N):
            napis = str( x[i][0] ) + ";" + str( x[i][1] ) + ";" + str( x[i][2] ) + "\n"
            plik.write(napis.replace(".",","))
        plik.close()
        raw_input(u'\nNaciśnij ENTER aby zakończyć...\n')
    

if __name__ == "__main__":
    rysuj()
